package com.allen.flink.dbus.sink;

import lombok.extern.slf4j.Slf4j;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HColumnDescriptor;
import org.apache.hadoop.hbase.HTableDescriptor;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.Connection;
import org.apache.hadoop.hbase.client.ConnectionFactory;
import org.apache.hadoop.hbase.client.Delete;
import org.apache.hadoop.hbase.client.HBaseAdmin;
import org.apache.hadoop.hbase.client.HTable;
import org.apache.hadoop.hbase.client.Put;
import org.apache.hadoop.hbase.util.Bytes;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;

/**
 * 功能: HBase 模板
 *
 * @date: 2020-03-26 16:53
 * @author: Allen
 * @version: 0.0.4-snapshot
 * @Email: allenZyhang@163.com
 * @since: JDK 1.8
 **/
@Slf4j
public class HBaseTemplate {
    private Configuration config; // hbase配置对象
    private Connection conn; // hbase连接
    
    public HBaseTemplate(final Configuration config) {
        this.config = config;
        initConn();
    }
    
    private void initConn() {
        try {
            this.conn = ConnectionFactory.createConnection(config);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
    
    public Connection getConn() {
        if (conn == null || conn.isClosed() || conn.isClosed()) {
            initConn();
        }
        return conn;
    }
    
    public boolean tableExists(String tableName) {
        try (HBaseAdmin admin = (HBaseAdmin) getConn().getAdmin()) {
            return admin.tableExists(TableName.valueOf(tableName));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
    
    public void createTable(String tableName, String... familyNames) {
        try (HBaseAdmin admin = (HBaseAdmin) getConn().getAdmin()) {
            
            HTableDescriptor tableDesc = new HTableDescriptor(TableName.valueOf(tableName));
            //添加列簇
            if (familyNames != null) {
                for (final String familyName : familyNames) {
                    HColumnDescriptor columnDesc = new HColumnDescriptor(familyName);
                    tableDesc.addFamily(columnDesc);
                }
            }
            admin.createTable(tableDesc);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
    
    public void disableTable(String tableName) {
        try (HBaseAdmin admin = (HBaseAdmin) getConn().getAdmin()) {
            admin.disableTable(tableName);
        } catch (IOException e) {
            log.error(e.getMessage(), e);
            throw new RuntimeException(e);
        }
    }
    
    public void deleteTable(String tableName) {
        try (HBaseAdmin admin = (HBaseAdmin) getConn().getAdmin()) {
            if (admin.isTableEnabled(tableName)) {
                disableTable(tableName);
            }
            admin.deleteTable(tableName);
        } catch (IOException e) {
            log.error(e.getMessage(), e);
            throw new RuntimeException(e);
        }
    }
    
    /**
     * 插入一行数据
     *
     * @param tableName 表名
     * @param row 行数据对象
     *
     * @return 是否成功
     */
    public boolean put(String tableName, HRow row) {
        boolean flag = false;
        try {
            HTable table = (HTable) getConn().getTable(TableName.valueOf(tableName));
            Put put = new Put(row.getRowKey());
            for (final HCell cell : row.getCells()) {
                put.addColumn(Bytes.toBytes(cell.getFamily()), Bytes.toBytes(cell.getQualifier()), cell.getValue());
            }
            table.put(put);
            flag = true;
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            throw new RuntimeException(e);
        }
        return flag;
    }
    
    /**
     * 批量插入
     *
     * @param tableName 表名
     * @param rows 行数据对象集合
     *
     * @return 是否成功
     */
    public boolean puts(String tableName, List<HRow> rows) {
        boolean flag = false;
        try {
            HTable table = (HTable) getConn().getTable(TableName.valueOf(tableName));
            List<Put> puts = new ArrayList<>();
            for (final HRow row : rows) {
                Put put = new Put(row.getRowKey());
                for (final HCell cell : row.getCells()) {
                    put.addColumn(Bytes.toBytes(cell.getFamily()), Bytes.toBytes(cell.getQualifier()), cell.getValue());
                }
                puts.add(put);
            }
            
            if (!puts.isEmpty()) {
                table.put(puts);
            }
            flag = true;
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            throw new RuntimeException(e);
        }
        return flag;
    }
    
    /**
     * 删除数据
     *
     * @param tableName 表名
     * @param rowKey rowKey
     *
     * @return 是否成功
     */
    public boolean delete(String tableName, byte[] rowKey) {
        boolean flag = false;
        try {
            HTable table = (HTable) getConn().getTable(TableName.valueOf(tableName));
            
            Delete delete = new Delete(rowKey);
            if (!delete.isEmpty()) {
                table.delete(delete);
            }
            flag = true;
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            throw new RuntimeException(e);
        }
        return flag;
    }
    
    /**
     * 批量删除数据
     *
     * @param tableName 表名
     * @param rowKeys rowKey集合
     *
     * @return 是否成功
     */
    public boolean deletes(String tableName, Set<byte[]> rowKeys) {
        boolean flag = false;
        try {
            HTable table = (HTable) getConn().getTable(TableName.valueOf(tableName));
            List<Delete> deletes = new ArrayList<>();
            for (final byte[] rowKey : rowKeys) {
                Delete delete = new Delete(rowKey);
                deletes.add(delete);
            }
            if (!deletes.isEmpty()) {
                table.delete(deletes);
            }
            flag = true;
        } catch (Exception e) {
            log.error(e.getMessage(), e);
            throw new RuntimeException(e);
        }
        return flag;
    }
    
    public void close() throws IOException {
        if (conn != null) {
            conn.close();
        }
    }
}
